import argparse
import subprocess
parser = argparse.ArgumentParser("Scheduling Task")
parser.add_argument('-t','--tasks', nargs='+', help='<Required> Set task list', required=True)
parser.add_argument('-s','--status', nargs='+', type=int, help='<Required> Set status', required=True)

args = parser.parse_args()

# python leader.py -l test1 test2 test2 test1 test2 test1 -s 1 1 1 1
gpu = args.status
p_list = []
ocu = {}
if __name__ == "__main__":
    while True:
        if len(args.tasks) == 0 and len(p_list) == 0:
            break

        try:    # 若有可用gpu, 状态1表示可用
            idx = gpu.index(1)
        except Exception as e:
            idx = None

        if idx is not None and len(args.tasks) > 0:
            job = args.tasks.pop()
            p_list.append(subprocess.Popen(f'./cmd/{job}.sh', shell=True))
            gpu[idx] = 0
            ocu[p_list[-1].pid] = idx
            print(ocu)
            # print(f'{job} pid: {p_list[-1].pid}')
            # print(list(map(lambda x: x.pid, p_list)))
            # print(gpu)
            print()
        for w in p_list:
            if w.poll() == 0:
                gpu[ocu[w.pid]] = 1
                ocu.pop(w.pid)
                p_list.remove(w)

        # p.returncode
        # p.wait()
